"""
区域匹配模块 - 专注于区域匹配和连通性分析

此模块负责匹配手指指向的内容区域，进行连通区域分析，
以及边界框的合并和处理。
"""
import cv2
import numpy as np
from typing import List, Tuple, Optional
from .content_analyzer import ContentRegion


class RegionMatcher:
    """
    区域匹配器类，专注于区域匹配和连通性分析
    
    功能包括：
    - 查找手指指向的内容区域
    - 连通区域分析
    - 边界框合并
    - 区域生长算法
    """
    
    def __init__(self, region_growing_iterations: int = 3):
        """
        初始化区域匹配器
        
        参数:
            region_growing_iterations: 区域生长迭代次数
        """
        self.region_growing_iterations = region_growing_iterations
    
    def find_pointed_content(self, content_regions: List[ContentRegion], 
                           finger_mask: np.ndarray,
                           min_overlap_ratio: float = 0.05) -> Optional[ContentRegion]:
        """
        查找手指指向的内容区域
        
        参数:
            content_regions: 内容区域列表
            finger_mask: 手指方向掩码
            min_overlap_ratio: 最小重叠比例
            
        返回:
            Optional[ContentRegion]: 指向的内容区域
        """
        if not content_regions or finger_mask is None:
            return None
            
        best_region = None
        max_overlap = 0
        
        for region in content_regions:
            # 计算区域掩码与手指掩码的重叠
            overlap = cv2.bitwise_and(region.binary_mask, finger_mask)
            overlap_count = np.count_nonzero(overlap)
            
            # 计算重叠比例
            region_count = np.count_nonzero(region.binary_mask)
            if region_count > 0:
                overlap_ratio = overlap_count / region_count
            else:
                overlap_ratio = 0
                
            # 更新最佳匹配区域
            if overlap_ratio >= min_overlap_ratio and overlap_count > max_overlap:
                max_overlap = overlap_count
                best_region = region
                
        return best_region
    
    def extract_connected_regions(self, binary_image: np.ndarray,
                                finger_mask: np.ndarray,
                                min_area: int = 100,
                                growth_iterations: Optional[int] = None) -> List[Tuple[int, int, int, int]]:
        """
        提取与手指指向区域相连的二值图像区域
        
        参数:
            binary_image: 二值化图像
            finger_mask: 手指方向掩码
            min_area: 最小区域面积
            growth_iterations: 区域生长迭代次数，如果为None则使用默认值
            
        返回:
            List[Tuple[int, int, int, int]]: 区域边界框列表 (x, y, w, h)
        """
        if binary_image is None or finger_mask is None:
            return []
        
        iterations = growth_iterations if growth_iterations is not None else self.region_growing_iterations
            
        # 初始种子点为手指掩码与二值图像的交集
        overlap = cv2.bitwise_and(binary_image, finger_mask)
        
        # 如果没有交集，扩大手指掩码后再试
        if np.count_nonzero(overlap) == 0:
            overlap = self._expand_and_retry(binary_image, finger_mask)
            
        # 如果仍然没有交集，返回手指掩码区域
        if np.count_nonzero(overlap) == 0:
            return self._get_finger_mask_boxes(finger_mask)
            
        # 进行连通区域分析和区域生长
        result_mask = self._perform_region_growing(binary_image, overlap, iterations)
        
        # 提取边界框
        boxes = self._extract_boxes_from_mask(result_mask, min_area)
        
        # 如果没有找到足够大的区域，返回备用区域
        if not boxes:
            boxes = self._get_fallback_boxes(binary_image, finger_mask)
                
        return boxes
    
    def _expand_and_retry(self, binary_image: np.ndarray, finger_mask: np.ndarray) -> np.ndarray:
        """扩大掩码后重试"""
        kernel = np.ones((5, 5), np.uint8)
        expanded_mask = cv2.dilate(finger_mask, kernel, iterations=2)
        return cv2.bitwise_and(binary_image, expanded_mask)
    
    def _get_finger_mask_boxes(self, finger_mask: np.ndarray) -> List[Tuple[int, int, int, int]]:
        """获取手指掩码的边界框"""
        rows, cols = np.where(finger_mask > 0)
        if len(rows) > 0 and len(cols) > 0:
            y_min, y_max = rows.min(), rows.max()
            x_min, x_max = cols.min(), cols.max()
            return [(x_min, y_min, x_max - x_min, y_max - y_min)]
        return []
    
    def _perform_region_growing(self, binary_image: np.ndarray, 
                              overlap: np.ndarray, iterations: int) -> np.ndarray:
        """执行区域生长"""
        # 标记所有连通区域
        retval, labels = cv2.connectedComponents(overlap)
        
        # 合并所有标记的区域
        result_mask = np.zeros_like(binary_image)
        
        for label in range(1, retval):  # 跳过背景标签0
            label_mask = (labels == label).astype(np.uint8) * 255
            
            # 使用标记区域作为种子，进行区域生长
            grown_mask = self._region_growing(binary_image, label_mask, iterations)
            
            # 合并到结果掩码
            result_mask = cv2.bitwise_or(result_mask, grown_mask)
            
        return result_mask
    
    def _extract_boxes_from_mask(self, mask: np.ndarray, min_area: int) -> List[Tuple[int, int, int, int]]:
        """从掩码中提取边界框"""
        contours, _ = cv2.findContours(
            mask, 
            cv2.RETR_EXTERNAL, 
            cv2.CHAIN_APPROX_SIMPLE
        )
        
        boxes = []
        for contour in contours:
            area = cv2.contourArea(contour)
            if area >= min_area:
                x, y, w, h = cv2.boundingRect(contour)
                boxes.append((x, y, w, h))
        
        return boxes
    
    def _get_fallback_boxes(self, binary_image: np.ndarray, 
                          finger_mask: np.ndarray) -> List[Tuple[int, int, int, int]]:
        """获取备用边界框"""
        # 计算手指方向掩码与二值图像的重叠区域的边界框
        combined = cv2.bitwise_and(binary_image, finger_mask)
        if np.count_nonzero(combined) > 0:
            rows, cols = np.where(combined > 0)
            y_min, y_max = rows.min(), rows.max()
            x_min, x_max = cols.min(), cols.max()
            return [(x_min, y_min, x_max - x_min, y_max - y_min)]
        
        # 如果仍然没有区域，但有手指掩码，返回掩码区域
        if np.count_nonzero(finger_mask) > 0:
            return self._get_finger_mask_boxes(finger_mask)
        
        return []
    
    def _region_growing(self, binary_image: np.ndarray, seed_mask: np.ndarray, 
                       iterations: int = 2) -> np.ndarray:
        """
        区域生长算法
        
        参数:
            binary_image: 二值化图像
            seed_mask: 种子掩码
            iterations: 生长迭代次数
            
        返回:
            np.ndarray: 生长后的区域掩码
        """
        # 形态学操作：先膨胀再腐蚀，填充小空洞
        kernel = np.ones((3, 3), np.uint8)
        result = seed_mask.copy()
        
        for _ in range(iterations):
            # 膨胀扩大区域
            dilated = cv2.dilate(result, kernel, iterations=1)
            
            # 找出新扩展的区域
            new_region = cv2.subtract(dilated, result)
            
            # 将新区域与二值图像进行与操作，只保留有内容的部分
            new_content = cv2.bitwise_and(binary_image, new_region)
            
            # 合并到结果中
            result = cv2.bitwise_or(result, new_content)
        
        # 最后再膨胀一次，连接相近区域
        final_result = cv2.dilate(result, kernel, iterations=1)
        
        return final_result
    
    def merge_overlapping_boxes(self, boxes: List[Tuple[int, int, int, int]], 
                              iou_threshold: float = 0.2) -> List[Tuple[int, int, int, int]]:
        """
        合并重叠的边界框
        
        参数:
            boxes: 边界框列表 (x, y, w, h)
            iou_threshold: IOU阈值
            
        返回:
            List[Tuple[int, int, int, int]]: 合并后的边界框列表
        """
        if not boxes:
            return []
            
        result_boxes = []
        remaining_boxes = list(boxes)
        
        while remaining_boxes:
            current_box = remaining_boxes.pop(0)
            i = 0
            
            while i < len(remaining_boxes):
                iou = self._calculate_iou(current_box, remaining_boxes[i])
                
                if iou > iou_threshold:
                    # 合并框
                    current_box = self._merge_boxes(current_box, remaining_boxes[i])
                    remaining_boxes.pop(i)
                else:
                    i += 1
            
            result_boxes.append(current_box)
            
        return result_boxes
    
    def _calculate_iou(self, box1: Tuple[int, int, int, int], 
                      box2: Tuple[int, int, int, int]) -> float:
        """计算两个边界框的IOU"""
        # 转换为(x1, y1, x2, y2)格式
        x1_1, y1_1, w1, h1 = box1
        x2_1, y2_1 = x1_1 + w1, y1_1 + h1
        
        x1_2, y1_2, w2, h2 = box2
        x2_2, y2_2 = x1_2 + w2, y1_2 + h2
        
        # 计算交集区域
        x1_i = max(x1_1, x1_2)
        y1_i = max(y1_1, y1_2)
        x2_i = min(x2_1, x2_2)
        y2_i = min(y2_1, y2_2)
        
        if x2_i <= x1_i or y2_i <= y1_i:
            return 0.0
        
        intersection = (x2_i - x1_i) * (y2_i - y1_i)
        
        # 计算各自面积
        area1 = w1 * h1
        area2 = w2 * h2
        
        # 计算IOU
        union = area1 + area2 - intersection
        
        return intersection / union if union > 0 else 0.0
    
    def _merge_boxes(self, box1: Tuple[int, int, int, int], 
                    box2: Tuple[int, int, int, int]) -> Tuple[int, int, int, int]:
        """合并两个边界框"""
        x1, y1, w1, h1 = box1
        x2, y2, w2, h2 = box2
        
        # 计算合并后的边界框
        x = min(x1, x2)
        y = min(y1, y2)
        w = max(x1 + w1, x2 + w2) - x
        h = max(y1 + h1, y2 + h2) - y
        
        return (x, y, w, h)
    
    def expand_box(self, box: Tuple[int, int, int, int], image_shape: Tuple[int, int], 
                  margin: int = 5) -> Tuple[int, int, int, int]:
        """
        扩展边界框
        
        参数:
            box: 边界框 (x, y, w, h)
            image_shape: 图像形状 (高度, 宽度)
            margin: 扩展边距
            
        返回:
            Tuple[int, int, int, int]: 扩展后的边界框
        """
        x, y, w, h = box
        h_img, w_img = image_shape
        
        # 扩展边界框
        x = max(0, x - margin)
        y = max(0, y - margin)
        w = min(w_img - x, w + margin * 2)
        h = min(h_img - y, h + margin * 2)
        
        return (x, y, w, h)
